#include "aclrtlaunch_add_custom.h"
#include "aclrtlaunch_sub_custom.h"

#include "torch_npu/csrc/core/npu/NPUStream.h"
#include <pybind11/pybind11.h>
#include <torch/extension.h>

namespace ops {
constexpr uint32_t FP16_SIZE = 2;
constexpr uint32_t BLOCK_SIZE = 16;

uint32_t AIV_CORE_NUM = 48;
uint32_t UB_SIZE = 196608;

inline __attribute__((always_inline)) uint32_t RoundUp(uint32_t num,
                                                       uint32_t rnd) {
  if (rnd == 0) {
    return 0;
  }
  return (num + rnd - 1) / rnd * rnd;
}

void add_custom(const at::Tensor &x, const at::Tensor &y, const at::Tensor &z) {
  uint32_t totalLengthX = at::numel(x);
  uint32_t totalLengthY = at::numel(y);

  uint32_t totalLengthA;
  uint32_t totalLengthB;
  void *aPtr;
  void *bPtr;
  if (totalLengthX > totalLengthY) {
    totalLengthA = totalLengthX;
    totalLengthB = totalLengthY;
    aPtr = x.data_ptr();
    bPtr = y.data_ptr();
  } else {
    totalLengthA = totalLengthY;
    totalLengthB = totalLengthX;
    aPtr = y.data_ptr();
    bPtr = x.data_ptr();
  }

  uint32_t coreNum = AIV_CORE_NUM;
  if (coreNum > totalLengthB) {
    coreNum = totalLengthB;
  }

  uint32_t pipe_buffer_num = 12;
  uint32_t maxTileLength = UB_SIZE / pipe_buffer_num / FP16_SIZE;

  uint32_t mode = 0; // at::ScalarType::Half
  if (x.scalar_type() == at::ScalarType::BFloat16) {
    mode = 1;
    maxTileLength = maxTileLength / 2;
  }

  uint32_t tileLength = totalLengthB / coreNum;
  if (tileLength > maxTileLength) {
    tileLength = maxTileLength;
  }
  tileLength = RoundUp(tileLength, BLOCK_SIZE);

  uint32_t totalTileNum = (totalLengthB + tileLength - 1) / tileLength;

  auto acl_stream = c10_npu::getCurrentNPUStream().stream();
  ACLRT_LAUNCH_KERNEL(add_custom)
  (coreNum, acl_stream, aPtr, bPtr, z.data_ptr(), totalLengthA, totalLengthB,
   totalTileNum, tileLength, mode);
}

void sub_custom(const at::Tensor &x, const at::Tensor &y, const at::Tensor &z) {
  uint32_t totalLengthX = at::numel(x);
  uint32_t totalLengthY = at::numel(y);

  uint32_t totalLengthA;
  uint32_t totalLengthB;
  void *aPtr;
  void *bPtr;
  if (totalLengthX > totalLengthY) {
    totalLengthA = totalLengthX;
    totalLengthB = totalLengthY;
    aPtr = x.data_ptr();
    bPtr = y.data_ptr();
  } else {
    totalLengthA = totalLengthY;
    totalLengthB = totalLengthX;
    aPtr = y.data_ptr();
    bPtr = x.data_ptr();
  }

  uint32_t coreNum = AIV_CORE_NUM;
  if (coreNum > totalLengthB) {
    coreNum = totalLengthB;
  }

  uint32_t pipe_buffer_num = 12;
  uint32_t maxTileLength = UB_SIZE / pipe_buffer_num / FP16_SIZE;

  uint32_t mode = 0; // at::ScalarType::Half
  if (x.scalar_type() == at::ScalarType::BFloat16) {
    mode = 1;
    maxTileLength = maxTileLength / 2;
  }

  uint32_t tileLength = totalLengthB / coreNum;
  if (tileLength > maxTileLength) {
    tileLength = maxTileLength;
  }
  tileLength = RoundUp(tileLength, BLOCK_SIZE);

  uint32_t totalTileNum = (totalLengthB + tileLength - 1) / tileLength;

  auto acl_stream = c10_npu::getCurrentNPUStream().stream();
  ACLRT_LAUNCH_KERNEL(sub_custom)
  (coreNum, acl_stream, aPtr, bPtr, z.data_ptr(), totalLengthA, totalLengthB,
   totalTileNum, tileLength, mode);
}

} // namespace ops

PYBIND11_MODULE(ascend_ops, m) {
  m.doc() = "ascend ops pybind11 interfaces";
  m.def("add", &ops::add_custom, "");
  m.def("sub", &ops::sub_custom, "");
}
